import argparse
import sys
from typing import List, Optional

class HexPatternSearcher:
    def __init__(self, filename: str):
        """初始化搜索器，加载文件"""
        self.filename = filename
        try:
            # 推荐使用 mmap 以支持大文件，对代码其他部分影响很小
            import mmap
            with open(filename, 'rb') as f:
                self.data = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
        except (FileNotFoundError, ValueError): # ValueError for empty file with mmap
            print(f"错误：文件 '{filename}' 不存在或为空。")
            sys.exit(1)
        except Exception as e:
            print(f"读取文件时出错：{e}")
            sys.exit(1)

    def __del__(self):
        """确保资源被释放，特别是 mmap 对象"""
        if hasattr(self, 'data') and self.data:
            self.data.close()

    def parse_pattern(self, pattern_str: str) -> List[Optional[int]]:
        """解析用户输入的模式字符串，如 '20 18 ? 46 36'"""
        parts = [p.strip() for p in pattern_str.split() if p.strip()]
        pattern = []
        for part in parts:
            if part == '?':
                pattern.append(None)
            else:
                try:
                    value = int(part, 16)
                    if not (0 <= value <= 255):
                        raise ValueError(f"字节值 {part} 超出范围 (0-255)")
                    pattern.append(value)
                except ValueError as e:
                    print(f"错误：无效的16进制值 '{part}'。请使用两位十六进制数（例如 4A, 0D）。\n详细信息: {e}")
                    sys.exit(1)
        return pattern

    def parse_deltas(self, deltas_str: str) -> List[Optional[int]]:
        """解析用户输入的增量规则字符串，如 '31 ? -70 195'"""
        parts = [p.strip() for p in deltas_str.split() if p.strip()]
        deltas = []
        for part in parts:
            if part == '?':
                deltas.append(None)
            else:
                try:
                    delta = int(part)
                    if not (-255 <= delta <= 255):
                        print(f"警告：增量值 {delta} 超出字节运算的可能范围 [-255, 255]，可能永远无法匹配。")
                    deltas.append(delta)
                except ValueError:
                    print(f"错误：无效的增量值 '{part}'。请输入整数或 '?'。")
                    sys.exit(1)
        return deltas

    def search_by_pattern(self, pattern: List[Optional[int]]) -> List[int]:
        """根据精确的字节模式搜索（支持 '?' 通配符）"""
        if not pattern:
            return []
        
        matches = []
        pattern_len = len(pattern)
        for i in range(len(self.data) - pattern_len + 1):
            match = True
            for j in range(pattern_len):
                if pattern[j] is not None and self.data[i + j] != pattern[j]:
                    match = False
                    break
            if match:
                matches.append(i)
        return matches

    def search_by_direct_deltas_previous(self, deltas: List[Optional[int]]) -> List[int]:
        """【新】使用直接输入的链式增量规则（支持 '?' 通配符）进行搜索。
        这是为了实现你的第3和第4个需求。
        """
        if not deltas:
            return []
        
        matches = []
        # 模式长度是增量数量 + 1
        pattern_len = len(deltas) + 1

        for i in range(len(self.data) - pattern_len + 1):
            match = True
            # last_known_byte_offset 是上一个已知字节相对于起始点i的偏移
            last_known_byte_offset = 0
            
            for j in range(len(deltas)):
                # 如果当前增量是 '?'，则跳过检查
                if deltas[j] is None:
                    continue

                # 当前字节的偏移
                current_byte_offset = j + 1
                
                # 从文件中获取上一个已知字节和当前字节
                prev_byte = self.data[i + last_known_byte_offset]
                current_byte = self.data[i + current_byte_offset]

                # 检查增量是否匹配
                if current_byte - prev_byte != deltas[j]:
                    match = False
                    break
                
                # 如果匹配，更新“上一个已知字节”的偏移为当前字节的偏移
                last_known_byte_offset = current_byte_offset

            if match:
                matches.append(i)
        
        return matches

    def search_by_direct_deltas_first(self, deltas: List[Optional[int]]) -> List[int]:
        """【新】使用直接输入的相对首位增量规则（支持 '?' 通配符）进行搜索。
        这是为了实现你的第5和第6个需求。
        """
        if not deltas:
            return []
        
        matches = []
        pattern_len = len(deltas) + 1
        
        for i in range(len(self.data) - pattern_len + 1):
            match = True
            first_byte = self.data[i]
            
            for j in range(len(deltas)):
                # 如果增量是 '?'，直接跳过
                if deltas[j] is None:
                    continue
                
                # 计算实际增量（相对于第一个字节）
                current_byte = self.data[i + j + 1]
                if current_byte - first_byte != deltas[j]:
                    match = False
                    break
            
            if match:
                matches.append(i)
        
        return matches

    def print_matches(self, matches: List[int], pattern_len: int, context_bytes: int = 8):
        """打印匹配位置和上下文（已修复标记位置的bug）"""
        if not matches:
            print("未找到匹配项")
            return
        
        print(f"\n找到 {len(matches)} 个匹配项：")
        for pos in matches:
            print("-" * 80)
            print(f"位置: 0x{pos:08X} ({pos})")
            
            start = max(0, pos - context_bytes)
            end = min(len(self.data), pos + pattern_len + context_bytes)
            
            # 打印上下文
            data_slice = self.data[start:end]
            for offset in range(0, len(data_slice), 16):
                chunk = data_slice[offset:offset+16]
                addr = start + offset
                hex_str = ' '.join(f'{b:02x}' for b in chunk)
                ascii_str = ''.join(chr(b) if 32 <= b <= 126 else '.' for b in chunk)
                print(f'{addr:08x}: {hex_str:<48} {ascii_str}')
            
            # 打印标记 (修正版)
            # 计算标记的起始行地址
            line_start_addr = pos - (pos % 16)
            # 计算标记在行内的缩进
            # 10 = "xxxxxxxx: "的长度
            # (pos % 16) * 3 = 在十六进制视图中的偏移
            indent = 10 + (pos % 16) * 3
            marker = ' ' * indent + '^' + '~' * (pattern_len * 3 - 1)
            print(marker)

def main():
    parser = argparse.ArgumentParser(
        description="""高级16进制文件搜索工具。

核心用法说明:
- 在 'pattern' 模式下，输入值为【十六进制】字节 (例如: 4A, FF, 0D)。
- 在 'delta-*' 模式下，输入值为【十进制】增量 (例如: 10, -1, 128)。
""",
        formatter_class=argparse.RawTextHelpFormatter  # 保持帮助信息中的换行和格式
    )
    
    parser.add_argument(
        "filename", 
        help="要搜索的文件名"
    )
    
    parser.add_argument(
        "pattern", 
        help="""要搜索的模式字符串。
根据所选模式，这里的数值应为十六进制或十进制。
支持 '?' 作为通配符。"""
    )
    
    parser.add_argument(
        "-m", "--mode", 
        choices=['pattern', 'delta-prev', 'delta-first'],
        default='pattern', 
        help="""指定搜索模式 (默认为: pattern):
  pattern:       精确/通配符模式。
                 输入为【十六进制】字节值。
                 示例: '40 00 ? ? 14'

  delta-prev:    链式增量模式 (字节与前一个字节比较)。
                 输入为【十进制】增量值。
                 示例: '-64 96' (表示B比A小64, C比B大96)

  delta-first:   相对首位增量模式 (所有字节与第一个字节比较)。
                 输入为【十进制】增量值。
                 示例: '-64 32' (表示B比A小64, C比A大32)
"""
    )
    
    parser.add_argument(
        "-c", "--context", 
        type=int, 
        default=16, 
        help="显示匹配位置前后的字节数 (默认: 16)"
    )
    
    args = parser.parse_args()
    
    searcher = HexPatternSearcher(args.filename)
    matches = []
    pattern_len = 0

    if args.mode == 'pattern':
        print("\n[模式: 精确/通配符搜索 (输入为十六进制)]")
        pattern = searcher.parse_pattern(args.pattern)
        pattern_len = len(pattern)
        print(f"解析的模式: {[f'{b:02x}' if b is not None else '?' for b in pattern]}")
        if pattern_len > 0:
            matches = searcher.search_by_pattern(pattern)
        
    elif args.mode == 'delta-prev':
        print("\n[模式: 链式增量搜索 (输入为十进制)]")
        deltas = searcher.parse_deltas(args.pattern)
        pattern_len = len(deltas) + 1
        print(f"解析的增量规则: {deltas}")
        if pattern_len > 1:
            matches = searcher.search_by_direct_deltas_previous(deltas)

    elif args.mode == 'delta-first':
        print("\n[模式: 相对首位增量搜索 (输入为十进制)]")
        deltas = searcher.parse_deltas(args.pattern)
        pattern_len = len(deltas) + 1
        print(f"解析的增量规则: {deltas}")
        if pattern_len > 1:
            matches = searcher.search_by_direct_deltas_first(deltas)

    if matches:
        searcher.print_matches(matches, pattern_len, args.context)
    else:
        print("\n未找到任何匹配项。")

if __name__ == "__main__":
    main()